{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Citations\n",
    "\n",
    "How can we get a model to cite which parts of the source documents it referenced in its response?\n",
    "\n",
    "To explore some techniques for extracting citations, let's first create a simple RAG chain. To start we'll just retrieve from the web using the [TavilySearchAPIRetriever](https://js.langchain.com/docs/integrations/retrievers/tavily)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "### Dependencies\n",
    "\n",
    "We’ll use an OpenAI chat model and embeddings and a Memory vector store in this walkthrough, but everything shown here works with any [ChatModel](/docs/modules/model_io/chat) or [LLM](/docs/modules/model_io/llms), [Embeddings](https://js.langchain.com/docs/modules/data_connection/text_embedding/), and [VectorStore](https://js.langchain.com/docs/modules/data_connection/vectorstores/) or [Retriever](/docs/modules/data_connection/retrievers/).\n",
    "\n",
    "We’ll use the following packages:\n",
    "\n",
    "```bash\n",
    "npm install --save langchain @langchain/community @langchain/openai\n",
    "```\n",
    "\n",
    "We need to set environment variables for Tavily Search & OpenAI:\n",
    "\n",
    "```bash\n",
    "export OPENAI_API_KEY=YOUR_KEY\n",
    "export TAVILY_API_KEY=YOUR_KEY\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LangSmith\n",
    "\n",
    "Many of the applications you build with LangChain will contain multiple steps with multiple invocations of LLM calls. As these applications get more and more complex, it becomes crucial to be able to inspect what exactly is going on inside your chain or agent. The best way to do this is with [LangSmith](https://smith.langchain.com/).\n",
    "\n",
    "Note that LangSmith is not needed, but it is helpful. If you do want to use LangSmith, after you sign up at the link above, make sure to set your environment variables to start logging traces:\n",
    "\n",
    "\n",
    "```bash\n",
    "export LANGCHAIN_TRACING_V2=true\n",
    "export LANGCHAIN_API_KEY=YOUR_KEY\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initial setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import { TavilySearchAPIRetriever } from \"@langchain/community/retrievers/tavily_search_api\";\n",
    "import { ChatPromptTemplate } from \"@langchain/core/prompts\";\n",
    "import { ChatOpenAI } from \"@langchain/openai\";\n",
    "\n",
    "const llm = new ChatOpenAI({\n",
    "  modelName: \"gpt-3.5-turbo\",\n",
    "  temperature: 0,\n",
    "});\n",
    "const retriever = new TavilySearchAPIRetriever({\n",
    "  k: 6,\n",
    "});\n",
    "const prompt = ChatPromptTemplate.fromMessages([\n",
    "  [\"system\", \"You're a helpful AI assistant. Given a user question and some web article snippets, answer the user question. If none of the articles answer the question, just say you don't know.\\n\\nHere are the web articles:{context}\"],\n",
    "  [\"human\", \"{question}\"],\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've got a model, retriever and prompt, let's chain them all together. We'll need to add some logic for formatting our retrieved `Document`s to a string that can be passed to our prompt. We'll make it so our chain returns both the answer and the retrieved Documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import { Document } from \"@langchain/core/documents\";\n",
    "import { StringOutputParser } from \"@langchain/core/output_parsers\";\n",
    "import { RunnableMap, RunnablePassthrough } from \"@langchain/core/runnables\";\n",
    "\n",
    "/**\n",
    " * Format the documents into a readable string.\n",
    " */\n",
    "const formatDocs = (input: Record<string, any>): string => {\n",
    "  const { docs } = input;\n",
    "  return \"\\n\\n\" + docs.map((doc: Document) => `Article title: ${doc.metadata.title}\\nArticle Snippet: ${doc.pageContent}`).join(\"\\n\\n\");\n",
    "}\n",
    "// subchain for generating an answer once we've done retrieval\n",
    "const answerChain = prompt.pipe(llm).pipe(new StringOutputParser());\n",
    "const map = RunnableMap.from({\n",
    "  question: new RunnablePassthrough(),\n",
    "  docs: retriever,\n",
    "})\n",
    "// complete chain that calls the retriever -> formats docs to string -> runs answer subchain -> returns just the answer and retrieved docs.\n",
    "const chain = map.assign({ context: formatDocs }).assign({ answer: answerChain }).pick([\"answer\", \"docs\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "  answer: \u001b[32m\"Cheetahs are capable of reaching speeds as high as 75 mph or 120 km/h. Their average speed, however,\"\u001b[39m... 29 more characters,\n",
       "  docs: [\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Now, their only hope lies in the hands of human conservationists, working tirelessly to save the che\"\u001b[39m... 880 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Are Cheetahs, and Other Fascinating Facts About the World's ...\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.discovermagazine.com/planet-earth/how-fast-are-cheetahs-and-other-fascinating-facts-abou\"\u001b[39m... 21 more characters,\n",
       "        score: \u001b[33m0.93715\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"If a lion comes along, the cheetah will abandon its catch -- it can't fight off a lion, and chances \"\u001b[39m... 911 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"What makes a cheetah run so fast? | HowStuffWorks\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://animals.howstuffworks.com/mammals/cheetah-speed.htm\"\u001b[39m,\n",
       "        score: \u001b[33m0.93412\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The science of cheetah speed\\n\"\u001b[39m +\n",
       "        \u001b[32m\"The cheetah (Acinonyx jubatus) is the fastest land animal on Earth, cap\"\u001b[39m... 738 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Can a Cheetah Run? - ThoughtCo\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.thoughtco.com/how-fast-can-a-cheetah-run-4587031\"\u001b[39m,\n",
       "        score: \u001b[33m0.93134\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"One of two videos from National Geographic's award-winning multimedia coverage of cheetahs in the ma\"\u001b[39m... 60 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"The Science of a Cheetah's Speed | National Geographic\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.youtube.com/watch?v=icFMTB0Pi0g\"\u001b[39m,\n",
       "        score: \u001b[33m0.93109\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Contact Us − +\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Address\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Smithsonian's National Zoo & Conservation Biology Institute  3001 Connecticut\"\u001b[39m... 1343 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah | Smithsonian's National Zoo and Conservation Biology Institute\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://nationalzoo.si.edu/animals/cheetah\"\u001b[39m,\n",
       "        score: \u001b[33m0.92938\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Threats to the Cheetah’s Reign\\n\"\u001b[39m +\n",
       "        \u001b[32m\"As unparalleled as the cheetah’s speed might be, they face numerous c\"\u001b[39m... 907 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Can a Cheetah Run? The Secrets Behind Its Incredible Speed\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.explorationjunkie.com/how-fast-can-a-cheetah-run/\"\u001b[39m,\n",
       "        score: \u001b[33m0.871\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    }\n",
       "  ]\n",
       "}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await chain.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangSmith trace [here](https://smith.langchain.com/public/bb0ed37e-b2be-4ae9-8b0d-ce2aff0b4b5e/r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function-calling\n",
    "\n",
    "### Cite documents\n",
    "Let's try using [OpenAI function-calling](/docs/modules/model_io/chat/function_calling) to make the model specify which of the provided documents it's actually referencing when answering. LangChain has some utils for converting objects or zod objects to the JSONSchema format expected by OpenAI, so we'll use that to define our functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import { z } from \"zod\";\n",
    "import { StructuredTool } from \"@langchain/core/tools\";\n",
    "import { formatToOpenAITool } from \"@langchain/openai\";\n",
    "\n",
    "class CitedAnswer extends StructuredTool {\n",
    "  name = \"cited_answer\";\n",
    "  \n",
    "  description = \"Answer the user question based only on the given sources, and cite the sources used.\";\n",
    "\n",
    "  schema = z.object({\n",
    "    answer: z.string().describe(\"The answer to the user question, which is based only on the given sources.\"),\n",
    "    citations: z.array(z.number()).describe(\"The integer IDs of the SPECIFIC sources which justify the answer.\")\n",
    "  });\n",
    "\n",
    "  constructor() {\n",
    "    super();\n",
    "  }\n",
    "\n",
    "  _call(input: z.infer<typeof this[\"schema\"]>): Promise<string> {\n",
    "    return Promise.resolve(JSON.stringify(input, null, 2));\n",
    "  }\n",
    "}\n",
    "\n",
    "const asOpenAITool = formatToOpenAITool(new CitedAnswer());\n",
    "const tools1 = [asOpenAITool];"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what the model output is like when we pass in our functions and a user input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage {\n",
       "  lc_serializable: \u001b[33mtrue\u001b[39m,\n",
       "  lc_kwargs: {\n",
       "    content: \u001b[32m\"\"\u001b[39m,\n",
       "    additional_kwargs: {\n",
       "      function_call: \u001b[90mundefined\u001b[39m,\n",
       "      tool_calls: [\n",
       "        {\n",
       "          id: \u001b[32m\"call_WzPoDCIRQ1pCah8k93cVrqex\"\u001b[39m,\n",
       "          type: \u001b[32m\"function\"\u001b[39m,\n",
       "          function: \u001b[36m[Object]\u001b[39m\n",
       "        }\n",
       "      ]\n",
       "    }\n",
       "  },\n",
       "  lc_namespace: [ \u001b[32m\"langchain_core\"\u001b[39m, \u001b[32m\"messages\"\u001b[39m ],\n",
       "  content: \u001b[32m\"\"\u001b[39m,\n",
       "  name: \u001b[90mundefined\u001b[39m,\n",
       "  additional_kwargs: {\n",
       "    function_call: \u001b[90mundefined\u001b[39m,\n",
       "    tool_calls: [\n",
       "      {\n",
       "        id: \u001b[32m\"call_WzPoDCIRQ1pCah8k93cVrqex\"\u001b[39m,\n",
       "        type: \u001b[32m\"function\"\u001b[39m,\n",
       "        function: {\n",
       "          name: \u001b[32m\"cited_answer\"\u001b[39m,\n",
       "          arguments: \u001b[32m\"{\\n\"\u001b[39m +\n",
       "            \u001b[32m`  \"answer\": \"Brian's height is 6'2\\\\\" - 3 inches\",\\n`\u001b[39m +\n",
       "            \u001b[32m'  \"citations\": [1, 3]\\n'\u001b[39m +\n",
       "            \u001b[32m\"}\"\u001b[39m\n",
       "        }\n",
       "      }\n",
       "    ]\n",
       "  }\n",
       "}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "const llmWithTool1 = llm.bind({\n",
    "  tools: tools1,\n",
    "  tool_choice: asOpenAITool\n",
    "});\n",
    "\n",
    "const exampleQ = `What Brian's height?\n",
    "\n",
    "Source: 1\n",
    "Information: Suzy is 6'2\"\n",
    "\n",
    "Source: 2\n",
    "Information: Jeremiah is blonde\n",
    "\n",
    "Source: 3\n",
    "Information: Brian is 3 inches shorted than Suzy`;\n",
    "\n",
    "await llmWithTool1.invoke(exampleQ);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangSmith trace [here](https://smith.langchain.com/public/34441213-cbb9-4775-a67e-2294aa1ccf69/r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll add an output parser to convert the OpenAI API response to a nice object. We use the [JsonOutputKeyToolsParser](https://api.js.langchain.com/classes/langchain_output_parsers.JsonOutputKeyToolsParser.html) for this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{ answer: \u001b[32m`Brian's height is 6'2\" - 3 inches`\u001b[39m, citations: [ \u001b[33m1\u001b[39m, \u001b[33m3\u001b[39m ] }"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import { JsonOutputKeyToolsParser } from \"langchain/output_parsers\";\n",
    "\n",
    "const outputParser = new JsonOutputKeyToolsParser({ keyName: \"cited_answer\", returnSingle: true });\n",
    "\n",
    "await llmWithTool1.pipe(outputParser).invoke(exampleQ);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangSmith trace [here](https://smith.langchain.com/public/1a045c25-ec5c-49f5-9756-6022edfea6af/r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we're ready to put together our chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import { Document } from \"@langchain/core/documents\";\n",
    "\n",
    "const formatDocsWithId = (docs: Array<Document>): string => {\n",
    "  return \"\\n\\n\" + docs.map((doc: Document, idx: number) => `Source ID: ${idx}\\nArticle title: ${doc.metadata.title}\\nArticle Snippet: ${doc.pageContent}`).join(\"\\n\\n\");\n",
    "}\n",
    "// subchain for generating an answer once we've done retrieval\n",
    "const answerChain1 = prompt.pipe(llmWithTool1).pipe(outputParser);\n",
    "const map1 = RunnableMap.from({\n",
    "  question: new RunnablePassthrough(),\n",
    "  docs: retriever,\n",
    "})\n",
    "// complete chain that calls the retriever -> formats docs to string -> runs answer subchain -> returns just the answer and retrieved docs.\n",
    "const chain1 = map1\n",
    "  .assign({ context: (input: { docs: Array<Document> }) => formatDocsWithId(input.docs) })\n",
    "  .assign({ cited_answer: answerChain1 })\n",
    "  .pick([\"cited_answer\", \"docs\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "  cited_answer: {\n",
       "    answer: \u001b[32m\"Cheetahs can reach speeds of up to 75 mph (120 km/h).\"\u001b[39m,\n",
       "    citations: [ \u001b[33m3\u001b[39m ]\n",
       "  },\n",
       "  docs: [\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The speeds attained by the cheetah may be only slightly greater than those achieved by the pronghorn\"\u001b[39m... 2527 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah - Wikipedia\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://en.wikipedia.org/wiki/Cheetah\"\u001b[39m,\n",
       "        score: \u001b[33m0.97773\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Contact Us − +\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Address\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Smithsonian's National Zoo & Conservation Biology Institute  3001 Connecticut\"\u001b[39m... 1343 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah | Smithsonian's National Zoo and Conservation Biology Institute\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://nationalzoo.si.edu/animals/cheetah\"\u001b[39m,\n",
       "        score: \u001b[33m0.9681\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The maximum speed cheetahs have been measured at is 114 km (71 miles) per hour, and they routinely r\"\u001b[39m... 1048 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah | Description, Speed, Habitat, Diet, Cubs, & Facts\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.britannica.com/animal/cheetah-mammal\"\u001b[39m,\n",
       "        score: \u001b[33m0.9459\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The science of cheetah speed\\n\"\u001b[39m +\n",
       "        \u001b[32m\"The cheetah (Acinonyx jubatus) is the fastest land animal on Earth, cap\"\u001b[39m... 738 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Can a Cheetah Run? - ThoughtCo\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.thoughtco.com/how-fast-can-a-cheetah-run-4587031\"\u001b[39m,\n",
       "        score: \u001b[33m0.93957\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"One of two videos from National Geographic's award-winning multimedia coverage of cheetahs in the ma\"\u001b[39m... 60 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"The Science of a Cheetah's Speed | National Geographic\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.youtube.com/watch?v=icFMTB0Pi0g\"\u001b[39m,\n",
       "        score: \u001b[33m0.92814\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"If a lion comes along, the cheetah will abandon its catch -- it can't fight off a lion, and chances \"\u001b[39m... 911 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"What makes a cheetah run so fast? | HowStuffWorks\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://animals.howstuffworks.com/mammals/cheetah-speed.htm\"\u001b[39m,\n",
       "        score: \u001b[33m0.85762\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    }\n",
       "  ]\n",
       "}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await chain1.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangSmith trace [here](https://smith.langchain.com/public/2a29cfd6-89fa-45bb-9b2a-f730e81061c2/r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cite snippets\n",
    "\n",
    "What if we want to cite actual text spans? We can try to get our model to return these, too.\n",
    "\n",
    "*Aside: Note that if we break up our documents so that we have many documents with only a sentence or two instead of a few long documents, citing documents becomes roughly equivalent to citing snippets, and may be easier for the model because the model just needs to return an identifier for each snippet instead of the actual text. Probably worth trying both approaches and evaluating.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "const citationSchema = z.object({\n",
    "  sourceId: z.number().describe(\"The integer ID of a SPECIFIC source which justifies the answer.\"),\n",
    "  quote: z.string().describe(\"The VERBATIM quote from the specified source that justifies the answer.\")\n",
    "})\n",
    "\n",
    "class QuotedAnswer extends StructuredTool {\n",
    "  name = \"quoted_answer\";\n",
    "  \n",
    "  description = \"Answer the user question based only on the given sources, and cite the sources used.\";\n",
    "\n",
    "  schema = z.object({\n",
    "    answer: z.string().describe(\"The answer to the user question, which is based only on the given sources.\"),\n",
    "    citations: z.array(citationSchema).describe(\"Citations from the given sources that justify the answer.\")\n",
    "  });\n",
    "\n",
    "  constructor() {\n",
    "    super();\n",
    "  }\n",
    "\n",
    "  _call(input: z.infer<typeof this[\"schema\"]>): Promise<string> {\n",
    "    return Promise.resolve(JSON.stringify(input, null, 2));\n",
    "  }\n",
    "}\n",
    "\n",
    "const quotedAnswerTool = formatToOpenAITool(new QuotedAnswer());\n",
    "const tools2 = [quotedAnswerTool];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import { Document } from \"@langchain/core/documents\";\n",
    "\n",
    "const outputParser2 = new JsonOutputKeyToolsParser({ keyName: \"quoted_answer\", returnSingle: true });\n",
    "const llmWithTool2 = llm.bind({\n",
    "  tools: tools2,\n",
    "  tool_choice: quotedAnswerTool,\n",
    "});\n",
    "const answerChain2 = prompt.pipe(llmWithTool2).pipe(outputParser2);\n",
    "const map2 = RunnableMap.from({\n",
    "  question: new RunnablePassthrough(),\n",
    "  docs: retriever,\n",
    "})\n",
    "// complete chain that calls the retriever -> formats docs to string -> runs answer subchain -> returns just the answer and retrieved docs.\n",
    "const chain2 = map2\n",
    "  .assign({ context: (input: { docs: Array<Document> }) => formatDocsWithId(input.docs) })\n",
    "  .assign({ quoted_answer: answerChain2 })\n",
    "  .pick([\"quoted_answer\", \"docs\"]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "  quoted_answer: {\n",
       "    answer: \u001b[32m\"Cheetahs can reach speeds of up to 70 mph.\"\u001b[39m,\n",
       "    citations: [\n",
       "      {\n",
       "        sourceId: \u001b[33m0\u001b[39m,\n",
       "        quote: \u001b[32m\"We’ve mentioned that these guys can reach speeds of up to 70 mph\"\u001b[39m\n",
       "      },\n",
       "      {\n",
       "        sourceId: \u001b[33m2\u001b[39m,\n",
       "        quote: \u001b[32m\"The maximum speed cheetahs have been measured at is 114 km (71 miles) per hour, and they routinely r\"\u001b[39m... 72 more characters\n",
       "      },\n",
       "      {\n",
       "        sourceId: \u001b[33m5\u001b[39m,\n",
       "        quote: \u001b[32m\"Cheetahs—the fastest land mammals on the planet—are able to reach speeds of up to 70 mph\"\u001b[39m\n",
       "      }\n",
       "    ]\n",
       "  },\n",
       "  docs: [\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"They are surprisingly graceful\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Cheetahs are very lithe-they move quickly and full-grown adults weigh\"\u001b[39m... 824 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Are Cheetahs - Proud Animal\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.proudanimal.com/2024/01/27/fast-cheetahs/\"\u001b[39m,\n",
       "        score: \u001b[33m0.97272\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The Science of Speed\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Instead, previous research has shown that the fastest animals are not the large\"\u001b[39m... 743 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Now Scientists Can Accurately Guess The Speed Of Any Animal\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.nationalgeographic.com/animals/article/Animal-speed-size-cheetahs\"\u001b[39m,\n",
       "        score: \u001b[33m0.96532\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The maximum speed cheetahs have been measured at is 114 km (71 miles) per hour, and they routinely r\"\u001b[39m... 1048 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah | Description, Speed, Habitat, Diet, Cubs, & Facts\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.britannica.com/animal/cheetah-mammal\"\u001b[39m,\n",
       "        score: \u001b[33m0.95122\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Now, their only hope lies in the hands of human conservationists, working tirelessly to save the che\"\u001b[39m... 880 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Are Cheetahs, and Other Fascinating Facts About the World's ...\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.discovermagazine.com/planet-earth/how-fast-are-cheetahs-and-other-fascinating-facts-abou\"\u001b[39m... 21 more characters,\n",
       "        score: \u001b[33m0.92667\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Contact Us − +\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Address\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Smithsonian's National Zoo & Conservation Biology Institute  3001 Connecticut\"\u001b[39m... 1343 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah | Smithsonian's National Zoo and Conservation Biology Institute\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://nationalzoo.si.edu/animals/cheetah\"\u001b[39m,\n",
       "        score: \u001b[33m0.91253\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Cheetahs—the fastest land mammals on the planet—are incredible creatures. They're able to reach spee\"\u001b[39m... 95 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Amazing Cheetah Facts | How Fast is a Cheetah? - Popular Mechanics\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.popularmechanics.com/science/animals/g30021998/facts-about-cheetahs/\"\u001b[39m,\n",
       "        score: \u001b[33m0.87489\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    }\n",
       "  ]\n",
       "}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await chain2.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangSmith trace [here](https://smith.langchain.com/public/2a032bc5-5b04-4dc3-8d85-49e5ec7e0157/r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Direct prompting\n",
    "\n",
    "Most models don't yet support function-calling. We can achieve similar results with direct prompting. Let's see what this looks like using an Anthropic chat model that is particularly proficient in working with XML:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup\n",
    "\n",
    "Install the LangChain Anthropic integration package:\n",
    "\n",
    "```bash\n",
    "npm install @langchain/anthropic\n",
    "```\n",
    "\n",
    "Add your Anthropic API key to your environment:\n",
    "\n",
    "```bash\n",
    "export ANTHROPIC_API_KEY=YOUR_KEY\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import { ChatAnthropic } from \"@langchain/anthropic\";\n",
    "import { ChatPromptTemplate } from \"@langchain/core/prompts\";\n",
    "\n",
    "const anthropic = new ChatAnthropic({\n",
    "  modelName: \"claude-instant-1.2\",\n",
    "});\n",
    "const system = `You're a helpful AI assistant. Given a user question and some web article snippets,\n",
    "answer the user question and provide citations. If none of the articles answer the question, just say you don't know.\n",
    "\n",
    "Remember, you must return both an answer and citations. A citation consists of a VERBATIM quote that\n",
    "justifies the answer and the ID of the quote article. Return a citation for every quote across all articles\n",
    "that justify the answer. Use the following format for your final output:\n",
    "\n",
    "<cited_answer>\n",
    "    <answer></answer>\n",
    "    <citations>\n",
    "        <citation><source_id></source_id><quote></quote></citation>\n",
    "        <citation><source_id></source_id><quote></quote></citation>\n",
    "        ...\n",
    "    </citations>\n",
    "</cited_answer>\n",
    "\n",
    "Here are the web articles:{context}`;\n",
    "\n",
    "const anthropicPrompt = ChatPromptTemplate.fromMessages([\n",
    "  [\"system\", system],\n",
    "  [\"human\", \"{question}\"]\n",
    "]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import { XMLOutputParser } from \"@langchain/core/output_parsers\";\n",
    "import { Document } from \"@langchain/core/documents\";\n",
    "import { RunnableLambda, RunnablePassthrough, RunnableMap } from \"@langchain/core/runnables\";\n",
    "\n",
    "const formatDocsToXML = (docs: Array<Document>): string => {\n",
    "  const formatted: Array<string> = [];\n",
    "  docs.forEach((doc, idx) => {\n",
    "    const docStr = `<source id=\"${idx}\">\n",
    "  <title>${doc.metadata.title}</title>\n",
    "  <article_snippet>${doc.pageContent}</article_snippet>\n",
    "</source>`\n",
    "    formatted.push(docStr);\n",
    "  });\n",
    "  return `\\n\\n<sources>${formatted.join(\"\\n\")}</sources>`;\n",
    "}\n",
    "\n",
    "const format3 = new RunnableLambda({\n",
    "  func: (input: { docs: Array<Document> }) => formatDocsToXML(input.docs)\n",
    "})\n",
    "const answerChain = anthropicPrompt\n",
    "  .pipe(anthropic)\n",
    "  .pipe(new XMLOutputParser())\n",
    "  .pipe(\n",
    "    new RunnableLambda({ func: (input: { cited_answer: any }) => input.cited_answer })\n",
    "  );\n",
    "const map3 = RunnableMap.from({\n",
    "  question: new RunnablePassthrough(),\n",
    "  docs: retriever,\n",
    "});\n",
    "const chain3 = map3.assign({ context: format3 }).assign({ cited_answer: answerChain }).pick([\"cited_answer\", \"docs\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "  cited_answer: [\n",
       "    {\n",
       "      answer: \u001b[32m\"Cheetahs can reach top speeds of between 60 to 70 mph.\"\u001b[39m\n",
       "    },\n",
       "    {\n",
       "      citations: [\n",
       "        { citation: \u001b[36m[Array]\u001b[39m },\n",
       "        { citation: \u001b[36m[Array]\u001b[39m },\n",
       "        { citation: \u001b[36m[Array]\u001b[39m }\n",
       "      ]\n",
       "    }\n",
       "  ],\n",
       "  docs: [\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"A cheetah's muscular tail helps control their steering and keep their balance when running very fast\"\u001b[39m... 210 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"75 Amazing Cheetah Facts Your Kids Will Love (2024)\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.mkewithkids.com/post/cheetah-facts-for-kids/\"\u001b[39m,\n",
       "        score: \u001b[33m0.97081\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The maximum speed cheetahs have been measured at is 114 km (71 miles) per hour, and they routinely r\"\u001b[39m... 1048 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah | Description, Speed, Habitat, Diet, Cubs, & Facts\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.britannica.com/animal/cheetah-mammal\"\u001b[39m,\n",
       "        score: \u001b[33m0.96824\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The Science of Speed\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Instead, previous research has shown that the fastest animals are not the large\"\u001b[39m... 743 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Now Scientists Can Accurately Guess The Speed Of Any Animal\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.nationalgeographic.com/animals/article/Animal-speed-size-cheetahs\"\u001b[39m,\n",
       "        score: \u001b[33m0.96237\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Contact Us − +\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Address\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Smithsonian's National Zoo & Conservation Biology Institute  3001 Connecticut\"\u001b[39m... 1343 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah | Smithsonian's National Zoo and Conservation Biology Institute\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://nationalzoo.si.edu/animals/cheetah\"\u001b[39m,\n",
       "        score: \u001b[33m0.94565\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"They are surprisingly graceful\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Cheetahs are very lithe-they move quickly and full-grown adults weigh\"\u001b[39m... 824 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Are Cheetahs - Proud Animal\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.proudanimal.com/2024/01/27/fast-cheetahs/\"\u001b[39m,\n",
       "        score: \u001b[33m0.91795\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Cheetahs are the world's fastest land animal. They can reach a speed of 69.5 miles per hour in just \"\u001b[39m... 100 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How fast is Tyreek Hill? 'The Cheetah' lives up to 40 time, Next Gen ...\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.sportingnews.com/us/nfl/news/fast-tyreek-hill-40-time-speed-chiefs/1cekgawhz39wr1tr472e4\"\u001b[39m... 5 more characters,\n",
       "        score: \u001b[33m0.83505\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    }\n",
       "  ]\n",
       "}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await chain3.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangSmith trace [here](https://smith.langchain.com/public/bebd86f5-ae9c-49ea-bc26-69c4fdf195b1/r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieval post-processing\n",
    "\n",
    "Another approach is to post-process our retrieved documents to compress the content, so that the source content is already minimal enough that we don't need the model to cite specific sources or spans. For example, we could break up each document into a sentence or two, embed those and keep only the most relevant ones. LangChain has some built-in components for this. Here we'll use a [RecursiveCharacterTextSplitter](https://js.langchain.com/docs/modules/data_connection/document_transformers/recursive_text_splitter), which creates chunks of a specified size by splitting on separator substrings, and an [EmbeddingsFilter](https://js.langchain.com/docs/modules/data_connection/retrievers/contextual_compression#embeddingsfilter), which keeps only the texts with the most relevant embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The maximum speed cheetahs have been measured at is 114 km (71 miles) per hour, and they routinely reach velocities of 80–100 km (50–62 miles) per hour while pursuing prey.\n",
      "cheetah,\n",
      "(Acinonyx jubatus), \n",
      "\n",
      "\n",
      "The science of cheetah speed\n",
      "The cheetah (Acinonyx jubatus) is the fastest land animal on Earth, capable of reaching speeds as high as 75 mph or 120 km/h. Cheetahs are predators that sneak up on their prey and sprint a short distance to chase and attack.\n",
      " Key Takeaways: How Fast Can a Cheetah Run?\n",
      "Fastest Cheetah on Earth \n",
      "\n",
      "\n",
      "Built for speed, the cheetah can accelerate from zero to 45 in just 2.5 seconds and reach top speeds of 60 to 70 mph, making it the fastest land mammal! Fun Facts\n",
      "Conservation Status\n",
      "Cheetah News\n",
      "Taxonomic Information\n",
      "Animal News\n",
      "NZCBI staff in Front Royal, Virginia, are mourning the loss of Walnut, a white-naped crane who became an internet sensation for choosing one of her keepers as her mate. \n",
      "\n",
      "\n",
      "Scientists calculate a cheetah's top speed is 75 mph, but the fastest recorded speed is somewhat slower. The top 10 fastest animals are: \n",
      "\n",
      "\n",
      "The pronghorn, an American animal resembling an antelope, is the fastest land animal in the Western Hemisphere. While a cheetah's top speed ranges from 65 to 75 mph (104 to 120 km/h), its average speed is only 40 mph (64 km/hr), punctuated by short bursts at its top speed. Basically, if a predator threatens to take a cheetah's kill or attack its young, a cheetah has to run. \n",
      "\n",
      "\n",
      "A cheetah eats a variety of small animals, including game birds, rabbits, small antelopes (including the springbok, impala, and gazelle), young warthogs, and larger antelopes (such as the kudu, hartebeest, oryx, and roan). Their faces are distinguished by prominent black lines that curve from the inner corner of each eye to the outer corners of the mouth, like a well-worn trail of inky tears. \n",
      "\n",
      "\n",
      "4 kg) Cheetah moms spend a lot of time teaching their cubs to chase, sometimes dragging live animals back to the den so the cubs can practice the chase-and-catch process \n",
      "\n",
      "\n",
      "Advertisement If confronted, a roughly 125-pound cheetah will always run rather than fight -- it's too weak, light and thin to have any chance against something like a lion, which can be twice as long as a cheetah and weigh more than 400 pounds (181 \n",
      "\n",
      "\n",
      "Cheetahs eat a variety of small animals, including game birds, rabbits, small antelopes (including the springbok, impala, and gazelle), young warthogs, and larger antelopes (such as the kudu, hartebeest, oryx, and roan) \n",
      "\n",
      "\n",
      "Historically, cheetahs ranged widely throughout Africa and Asia, from the Cape of Good Hope to the Mediterranean, throughout the Arabian Peninsula and the Middle East, from Israel, India and Pakistan north to the northern shores of the Caspian and Aral Seas, and west through Uzbekistan, Turkmenistan, Afghanistan, and Pakistan into central India. Header Links \n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import { RecursiveCharacterTextSplitter } from \"langchain/text_splitter\";\n",
    "import { EmbeddingsFilter } from \"langchain/retrievers/document_compressors/embeddings_filter\";\n",
    "import { OpenAIEmbeddings } from \"@langchain/openai\";\n",
    "import { DocumentInterface } from \"@langchain/core/documents\";\n",
    "import { RunnableMap, RunnablePassthrough } from \"@langchain/core/runnables\";\n",
    "\n",
    "const splitter = new RecursiveCharacterTextSplitter({\n",
    "  chunkSize: 400,\n",
    "  chunkOverlap: 0,\n",
    "  separators: [\"\\n\\n\", \"\\n\", \".\", \" \"],\n",
    "  keepSeparator: false,\n",
    "});\n",
    "\n",
    "const compressor = new EmbeddingsFilter({\n",
    "  embeddings: new OpenAIEmbeddings(),\n",
    "  k: 10,\n",
    "});\n",
    "\n",
    "const splitAndFilter = async (input): Promise<Array<DocumentInterface>> => {\n",
    "  const { docs, question } = input;\n",
    "  const splitDocs = await splitter.splitDocuments(docs);\n",
    "  const statefulDocs = await compressor.compressDocuments(splitDocs, question);\n",
    "  return statefulDocs;\n",
    "};\n",
    "\n",
    "const retrieveMap = RunnableMap.from({\n",
    "  question: new RunnablePassthrough(),\n",
    "  docs: retriever,\n",
    "});\n",
    "\n",
    "const retrieve = retrieveMap.pipe(splitAndFilter);\n",
    "const docs = await retrieve.invoke(\"How fast are cheetahs?\");\n",
    "for (const doc of docs) {\n",
    "  console.log(doc.pageContent, \"\\n\\n\");\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangSmith trace [here](https://smith.langchain.com/public/1bb61806-7d09-463d-909a-a7da410e79d4/r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "const chain4 = retrieveMap\n",
    "  .assign({ context: formatDocs })\n",
    "  .assign({ answer: answerChain })\n",
    "  .pick([\"answer\", \"docs\"]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "  answer: [\n",
       "    {\n",
       "      answer: \u001b[32m\"\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Cheetahs are the fastest land animals. They can reach top speeds of around 75 mph (120 km/h) and ro\"\u001b[39m... 74 more characters\n",
       "    },\n",
       "    { citations: [ { citation: \u001b[36m[Array]\u001b[39m }, { citation: \u001b[36m[Array]\u001b[39m } ] }\n",
       "  ],\n",
       "  docs: [\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The maximum speed cheetahs have been measured at is 114 km (71 miles) per hour, and they routinely r\"\u001b[39m... 1048 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"cheetah - Encyclopedia Britannica | Britannica\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.britannica.com/animal/cheetah-mammal\"\u001b[39m,\n",
       "        score: \u001b[33m0.97059\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Contact Us − +\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Address\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Smithsonian's National Zoo & Conservation Biology Institute  3001 Connecticut\"\u001b[39m... 1343 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://nationalzoo.si.edu/animals/cheetah\"\u001b[39m,\n",
       "        score: \u001b[33m0.95102\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The science of cheetah speed\\n\"\u001b[39m +\n",
       "        \u001b[32m\"The cheetah (Acinonyx jubatus) is the fastest land animal on Earth, cap\"\u001b[39m... 738 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Can a Cheetah Run?\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.thoughtco.com/how-fast-can-a-cheetah-run-4587031\"\u001b[39m,\n",
       "        score: \u001b[33m0.94974\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Now, their only hope lies in the hands of human conservationists, working tirelessly to save the che\"\u001b[39m... 880 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Are Cheetahs, and Other Fascinating Facts About the World's ...\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.discovermagazine.com/planet-earth/how-fast-are-cheetahs-and-other-fascinating-facts-abou\"\u001b[39m... 21 more characters,\n",
       "        score: \u001b[33m0.92695\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"One of two videos from National Geographic's award-winning multimedia coverage of cheetahs in the ma\"\u001b[39m... 60 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"The Science of a Cheetah's Speed | National Geographic\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.youtube.com/watch?v=icFMTB0Pi0g\"\u001b[39m,\n",
       "        score: \u001b[33m0.90754\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The speeds attained by the cheetah may be only slightly greater than those achieved by the pronghorn\"\u001b[39m... 2527 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah - Wikipedia\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://en.wikipedia.org/wiki/Cheetah\"\u001b[39m,\n",
       "        score: \u001b[33m0.89476\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    }\n",
       "  ]\n",
       "}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "// Note the documents have an article \"summary\" in the metadata that is now much longer than the\n",
    "// actual document page content. This summary isn't actually passed to the model.\n",
    "await chain4.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangSmith trace [here](https://smith.langchain.com/public/f93302e6-a31b-454e-9fc7-94fb4a931a9d/r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generation post-processing\n",
    "\n",
    "Another approach is to post-process our model generation. In this example we'll first generate just an answer, and then we'll ask the model to annotate it's own answer with citations. The downside of this approach is of course that it is slower and more expensive, because two model calls need to be made.\n",
    "\n",
    "Let's apply this to our initial chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import { StructuredTool } from \"@langchain/core/tools\";\n",
    "import { formatToOpenAITool } from \"@langchain/openai\";\n",
    "import { z } from \"zod\";\n",
    "\n",
    "class AnnotatedAnswer extends StructuredTool {\n",
    "  name = \"annotated_answer\";\n",
    "\n",
    "  description = \"Annotate the answer to the user question with quote citations that justify the answer\";\n",
    "\n",
    "  schema = z.object({\n",
    "    citations: z.array(citationSchema).describe(\"Citations from the given sources that justify the answer.\"),\n",
    "  })\n",
    "\n",
    "  _call(input: z.infer<typeof this[\"schema\"]>): Promise<string> {\n",
    "    return Promise.resolve(JSON.stringify(input, null, 2));\n",
    "  }\n",
    "}\n",
    "\n",
    "const annotatedAnswerTool = formatToOpenAITool(new AnnotatedAnswer());\n",
    "\n",
    "const llmWithTools5 = llm.bind({\n",
    "  tools: [annotatedAnswerTool],\n",
    "  tool_choice: annotatedAnswerTool,\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "import { ChatPromptTemplate, MessagesPlaceholder } from \"@langchain/core/prompts\";\n",
    "import { RunnableSequence } from \"@langchain/core/runnables\";\n",
    "import { JsonOutputKeyToolsParser } from \"langchain/output_parsers\";\n",
    "import { RunnableMap, RunnablePassthrough } from \"@langchain/core/runnables\";\n",
    "import { AIMessage, ToolMessage } from \"@langchain/core/messages\";\n",
    "\n",
    "const prompt5 = ChatPromptTemplate.fromMessages([\n",
    "  [\"system\", \"You're a helpful AI assistant. Given a user question and some web article snippets, answer the user question. If none of the articles answer the question, just say you don't know.\\n\\nHere are the web articles:{context}\"],\n",
    "  [\"human\", \"{question}\"],\n",
    "  new MessagesPlaceholder({\n",
    "    variableName: \"chat_history\",\n",
    "    optional: true,\n",
    "  }),\n",
    "  new MessagesPlaceholder({\n",
    "    variableName: \"toolMessage\",\n",
    "    optional: true,\n",
    "  })\n",
    "]);\n",
    "\n",
    "const answerChain5 = prompt5.pipe(llmWithTools5);\n",
    "const annotationChain = RunnableSequence.from([\n",
    "  prompt5,\n",
    "  llmWithTools5,\n",
    "  new JsonOutputKeyToolsParser({ keyName: \"annotated_answer\", returnSingle: true }),\n",
    "  (input: any) => input.citations\n",
    "]);\n",
    "const map5 = RunnableMap.from({\n",
    "  question: new RunnablePassthrough(),\n",
    "  docs: retriever,\n",
    "});\n",
    "const chain5 = map5\n",
    "  .assign({ context: formatDocs })\n",
    "  .assign({ aiMessage: answerChain5 })\n",
    "  .assign({\n",
    "    chat_history: (input) => input.aiMessage,\n",
    "    toolMessage: (input) => new ToolMessage({\n",
    "      tool_call_id: input.aiMessage.additional_kwargs.tool_calls[0].id,\n",
    "      content: input.aiMessage.additional_kwargs.content ?? \"\",\n",
    "    })\n",
    "  })\n",
    "  .assign({\n",
    "    annotations: annotationChain,\n",
    "  })\n",
    "  .pick([\"answer\", \"docs\", \"annotations\"]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "  docs: [\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"They are surprisingly graceful\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Cheetahs are very lithe-they move quickly and full-grown adults weigh\"\u001b[39m... 824 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Are Cheetahs - Proud Animal\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.proudanimal.com/2024/01/27/fast-cheetahs/\"\u001b[39m,\n",
       "        score: \u001b[33m0.96021\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Contact Us − +\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Address\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Smithsonian's National Zoo & Conservation Biology Institute  3001 Connecticut\"\u001b[39m... 1343 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah | Smithsonian's National Zoo and Conservation Biology Institute\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://nationalzoo.si.edu/animals/cheetah\"\u001b[39m,\n",
       "        score: \u001b[33m0.94798\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The science of cheetah speed\\n\"\u001b[39m +\n",
       "        \u001b[32m\"The cheetah (Acinonyx jubatus) is the fastest land animal on Earth, cap\"\u001b[39m... 738 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Can a Cheetah Run? - ThoughtCo\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.thoughtco.com/how-fast-can-a-cheetah-run-4587031\"\u001b[39m,\n",
       "        score: \u001b[33m0.92591\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The maximum speed cheetahs have been measured at is 114 km (71 miles) per hour, and they routinely r\"\u001b[39m... 1048 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Cheetah | Description, Speed, Habitat, Diet, Cubs, & Facts\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.britannica.com/animal/cheetah-mammal\"\u001b[39m,\n",
       "        score: \u001b[33m0.90128\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"The Science of Speed\\n\"\u001b[39m +\n",
       "        \u001b[32m\"Instead, previous research has shown that the fastest animals are not the large\"\u001b[39m... 743 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"Now Scientists Can Accurately Guess The Speed Of Any Animal\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.nationalgeographic.com/animals/article/Animal-speed-size-cheetahs\"\u001b[39m,\n",
       "        score: \u001b[33m0.90097\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    },\n",
       "    Document {\n",
       "      pageContent: \u001b[32m\"Now, their only hope lies in the hands of human conservationists, working tirelessly to save the che\"\u001b[39m... 880 more characters,\n",
       "      metadata: {\n",
       "        title: \u001b[32m\"How Fast Are Cheetahs, and Other Fascinating Facts About the World's ...\"\u001b[39m,\n",
       "        source: \u001b[32m\"https://www.discovermagazine.com/planet-earth/how-fast-are-cheetahs-and-other-fascinating-facts-abou\"\u001b[39m... 21 more characters,\n",
       "        score: \u001b[33m0.89788\u001b[39m,\n",
       "        images: \u001b[1mnull\u001b[22m\n",
       "      }\n",
       "    }\n",
       "  ],\n",
       "  annotations: [\n",
       "    {\n",
       "      sourceId: \u001b[33m0\u001b[39m,\n",
       "      quote: \u001b[32m\"We’ve mentioned that these guys can reach speeds of up to 70 mph, but did you know they can go from \"\u001b[39m... 22 more characters\n",
       "    },\n",
       "    {\n",
       "      sourceId: \u001b[33m1\u001b[39m,\n",
       "      quote: \u001b[32m\"Built for speed, the cheetah can accelerate from zero to 45 in just 2.5 seconds and reach top speeds\"\u001b[39m... 52 more characters\n",
       "    },\n",
       "    {\n",
       "      sourceId: \u001b[33m2\u001b[39m,\n",
       "      quote: \u001b[32m\"The maximum speed cheetahs have been measured at is 114 km (71 miles) per hour, and they routinely r\"\u001b[39m... 72 more characters\n",
       "    }\n",
       "  ]\n",
       "}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await chain5.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangSmith trace [here](https://smith.langchain.com/public/f4ca647d-b43d-49ba-8df5-65a9761f712e/r)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Deno",
   "language": "typescript",
   "name": "deno"
  },
  "language_info": {
   "file_extension": ".ts",
   "mimetype": "text/x.typescript",
   "name": "typescript",
   "nb_converter": "script",
   "pygments_lexer": "typescript",
   "version": "5.3.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
